# 一个简单的摄像头应用实例，利用了OpenCV打开摄像头
# 导入模型，类别文件，路径根据模型实际进行修改
# 打印预测结果，并在图片显示窗口打印出类别名称
# opencv 输出不了中文，所以中文字会变成问号

import cv2
from prediction import Prediction
from PIL import Image
import time


cap = cv2.VideoCapture("/dev/video0")
model_path = "/old_hard_disk/home/ailab/cat-vs-dog/model/model_best.pth.tar"
classs_path = "/old_hard_disk/home/ailab/cat-vs-dog/model/classes.txt"
pre = Prediction()
pre.load_retrained_model(model_path, classs_path)
#pre.load_pretrained_model("googlenet", "imagenet_classes.txt")

while True:
    t1 = time.perf_counter()
    ret, frame = cap.read()

    predict_image = cv2.resize(frame, (320, 240))
    predict_image = Image.fromarray(cv2.cvtColor(predict_image, cv2.COLOR_BGR2RGB))
    result = pre.predict(predict_image)
    if result:
        print(result)
        text = result[0][0] + " : " + (str(round(result[0][1], 2))) + "%"
        print(text)
        cv2.putText(frame, text, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.imshow("Test window", frame)
    print("interval:{}".format(time.perf_counter() - t1))
    if cv2.waitKey(1) == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()
